# -*- coding: utf-8 -*-

# created by mengqi on 2023/5/25

import cv2

MIN_W = 90
MIN_H = 90
# 检测线的高度
LINE_HEIGHT = 550
# 线的偏移
LINE_OFFSET = 7
# 统计车的数量
CARS_COUNT = 0
# 存放有效车辆的数组
CARS_ARR = []


def center(px, py, pw, ph):
    x1 = int(pw / 2)
    y1 = int(ph / 2)
    cx = px + x1
    cy = py + y1

    return cx, cy


CAP = cv2.VideoCapture('video.mp4')
BG_MOG = cv2.bgsegm.createBackgroundSubtractorMOG()

# 形态学kernel
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))

while True:
    ret, frame = CAP.read()
    if ret:
        # 灰度
        gray_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        # 去噪（高斯）
        blur = cv2.GaussianBlur(gray_img, (3, 3), 5)
        # 去背影
        mask = BG_MOG.apply(blur)
        # 腐蚀， 去掉图中小斑块
        erode = cv2.erode(mask, kernel)
        # 膨胀， 还原放大
        dilate = cv2.dilate(erode, kernel, iterations=3)
        # 闭操作，去掉物体内部的小块
        close = cv2.morphologyEx(dilate, cv2.MORPH_CLOSE, kernel)
        close = cv2.morphologyEx(close, cv2.MORPH_CLOSE, kernel)
        cnts, h = cv2.findContours(close, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # 画一条检测线
        cv2.line(frame, (10, LINE_HEIGHT), (1200, LINE_HEIGHT), (255, 255, 0), 3)

        for (i, c) in enumerate(cnts):
            (x, y, w, h) = cv2.boundingRect(c)
            # 对车辆的宽高进行判断
            # 以验证是否是有效的车辆
            isValid = (w >= MIN_W) and (h >= MIN_H)
            if not isValid:
                continue

            # 到这里都是有效的车
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
            c_point = center(x, y, w, h)
            CARS_ARR.append(c_point)
            cv2.circle(frame, c_point, 5, (0, 0, 255), -1)

            # 逻辑判断过线，但是没有目标跟踪，所以有误差，需要深度学习完成 by mengqi
            for (x, y) in CARS_ARR:
                if (y > LINE_HEIGHT - LINE_OFFSET) and (y < LINE_HEIGHT + LINE_OFFSET):
                    CARS_COUNT += 1
                    CARS_ARR.remove((x, y))
                    # print(car_count)

        cv2.putText(frame, "Cars Count:" + str(CARS_COUNT), (500, 60), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 5)
        cv2.imshow('video', frame)
        # cv2.imshow('erode', close)

    key = cv2.waitKey(1)
    if key > 0:
        break

CAP.release()
cv2.destroyAllWindows()
